use crate::common::packages::{
    ERROR_DATA_CHANGED, ERROR_DATA_LEN, FromSized, GetAction, PackageMetadata, Serializable,
    SetGetRole, merge, split,
};
use anyhow::{Error, Result};
use quinn::{
    Endpoint, ReadError, ReadExactError, RecvStream, SendStream, StoppedError, WriteError,
};
use quinn_proto::crypto::rustls::{QuicClientConfig, QuicServerConfig};
use quinn_proto::{ConnectionError, VarInt};
use rustls::RootCertStore;
use rustls::pki_types::pem::PemObject;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls::server::WebPkiClientVerifier;
use std::borrow::Cow;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::{
    Arc,
    atomic::{self, AtomicI8, AtomicI64},
};
use std::task::{Context, Poll};
use std::thread;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
use tokio::sync::RwLock;

pub fn init() {
    rustls::crypto::CryptoProvider::install_default(rustls::crypto::ring::default_provider())
        .unwrap()
}

pub struct QuicTls {
    ca_pem: Vec<u8>,
    cert_pem: Vec<u8>,
    key_pem: Vec<u8>,
}

impl QuicTls {
    pub fn new(ca_pem: &[u8], cert_pem: &[u8], key_pem: &[u8]) -> Self {
        QuicTls {
            ca_pem: ca_pem.to_vec(),
            cert_pem: cert_pem.to_vec(),
            key_pem: key_pem.to_vec(),
        }
    }

    pub fn client(&self) -> quinn::ClientConfig {
        let ca = CertificateDer::from(
            openssl::x509::X509::from_pem(&self.ca_pem)
                .unwrap()
                .to_der()
                .unwrap(),
        );
        let cert = CertificateDer::from(
            openssl::x509::X509::from_pem(&self.cert_pem)
                .unwrap()
                .to_der()
                .unwrap(),
        );
        let key = PrivateKeyDer::from_pem_slice(self.key_pem.as_slice()).unwrap();
        let mut root = RootCertStore::empty();
        root.add(ca).unwrap();

        let mut client_crypto = rustls::ClientConfig::builder()
            .with_root_certificates(Arc::new(root))
            .with_client_auth_cert(vec![cert], key)
            .unwrap();
        client_crypto.alpn_protocols = vec![b"hq-29".to_vec()];
        quinn::ClientConfig::new(Arc::new(QuicClientConfig::try_from(client_crypto).unwrap()))
    }

    pub fn server(&self) -> quinn::ServerConfig {
        let ca = CertificateDer::from(
            openssl::x509::X509::from_pem(&self.ca_pem)
                .unwrap()
                .to_der()
                .unwrap(),
        );
        let cert = CertificateDer::from(
            openssl::x509::X509::from_pem(&self.cert_pem)
                .unwrap()
                .to_der()
                .unwrap(),
        );
        let key = PrivateKeyDer::from_pem_slice(self.key_pem.as_slice()).unwrap();
        let mut root = RootCertStore::empty();
        root.add(ca).unwrap();
        let filter = WebPkiClientVerifier::builder(Arc::new(root))
            .build()
            .unwrap();
        let mut server_crypto = rustls::ServerConfig::builder()
            .with_client_cert_verifier(filter)
            .with_single_cert(vec![cert], key)
            .unwrap();
        server_crypto.alpn_protocols = vec![b"hq-29".to_vec()];
        let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(
            QuicServerConfig::try_from(server_crypto).unwrap(),
        ));
        let transport_cfg = Arc::get_mut(&mut server_config.transport).unwrap();
        transport_cfg.max_concurrent_uni_streams(0_u8.into());
        server_config
    }
}

pub async fn quic_server(tls: &QuicTls, endpoint: SocketAddr) -> Endpoint {
    Endpoint::server(tls.server(), endpoint).unwrap()
}

pub async fn quic_client(tls: &QuicTls) -> Endpoint {
    let mut e = Endpoint::client((std::net::Ipv4Addr::UNSPECIFIED, 0).into()).unwrap();
    e.set_default_client_config(tls.client());
    e
}

#[tokio::test]
async fn quic_demo() {
    init();

    let server_cert = include_bytes!("../ssl/server.crt");
    let server_key = include_bytes!("../ssl/server.key");
    let cli_cert = include_bytes!("../ssl/client.crt");
    let cli_key = include_bytes!("../ssl/client.key");
    let ca_pem = include_bytes!("../ssl/ca.crt");

    let server_tls = QuicTls::new(ca_pem, server_cert, server_key);
    let client_tls = QuicTls::new(ca_pem, cli_cert, cli_key);

    let server_end = quic_server(&server_tls, "0.0.0.0:8888".parse().unwrap()).await;
    tokio::spawn(async move {
        while let Some(conn) = server_end.accept().await {
            let ss = conn.await.unwrap();
            loop {
                let mut stream = ss.accept_bi().await.unwrap();
                let mut r = vec![0u8; 1024];
                let len = stream.1.read(&mut r).await.unwrap().unwrap();
                let recv = String::from_utf8_lossy(&r[..len]).to_string();
                println!("server receive {}: {:?}", len, recv);
                stream.0.write_all(recv.as_bytes()).await.unwrap();
                stream.0.flush().await.unwrap();
                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
            }
        }
    });

    let client_end = quic_client(&client_tls).await;
    let conn = client_end
        .connect("127.0.0.1:8888".parse().unwrap(), "www.li-ming.top")
        .unwrap()
        .await
        .unwrap();
    let client = tokio::spawn(async move {
        for _i in 0..3 {
            let mut stream = conn.open_bi().await.unwrap();
            stream.0.write_all(b"hello world").await.unwrap();
            stream.0.flush().await.unwrap();
            let mut buf = vec![0u8; 1024];
            let s = stream.1.read(&mut buf).await.unwrap().unwrap();
            println!(
                "client receive {} : {:?}",
                s,
                String::from_utf8_lossy(&buf[..s]).to_string()
            );
        }
    });
    client.await.unwrap();
}

pub struct QuicConn {
    conn: quinn::Connection,
    stream_num: Arc<AtomicI64>,
}

impl QuicConn {
    pub fn conn_num(&self) -> i64 {
        self.stream_num.load(atomic::Ordering::SeqCst)
    }

    pub fn peer_addr(&self) -> SocketAddr {
        self.conn.remote_address()
    }

    pub fn is_open(&self) -> bool {
        self.conn.close_reason().is_none()
    }

    pub fn new(conn: quinn::Connection) -> Arc<RwLock<Self>> {
        Arc::new(RwLock::new(QuicConn {
            conn,
            stream_num: Arc::new(AtomicI64::new(0)),
        }))
    }

    pub fn close(&self, code: u32, reason: &str) {
        self.conn.close(VarInt::from_u32(code), reason.as_bytes());
    }

    pub async fn accept_bi(conn: Arc<RwLock<QuicConn>>) -> Result<QuicStream> {
        let c = conn.write().await;
        match c.conn.accept_bi().await {
            Ok((a, b)) => Ok(QuicStream {
                conn: conn.clone(),
                send: a,
                recv: b,
                stat: AtomicI8::new(0),
            }),
            Err(e) => Err(handle_connection_error(&e)),
        }
    }

    pub async fn open_bi(conn: Arc<RwLock<QuicConn>>) -> Result<QuicStream> {
        let c = conn.write().await;
        match c.conn.open_bi().await {
            Ok((a, b)) => Ok(QuicStream {
                conn: conn.clone(),
                send: a,
                recv: b,
                stat: AtomicI8::new(0),
            }),
            Err(e) => Err(handle_connection_error(&e)),
        }
    }
}

pub struct QuicStream {
    conn: Arc<RwLock<QuicConn>>,
    send: SendStream,
    recv: RecvStream,
    stat: AtomicI8, // 流的状态，最低一位表示发送流，第二位表示接受流，置1表示流正常
}

impl AsyncRead for QuicStream {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.get_mut().recv).poll_read(cx, buf)
    }
}

impl AsyncWrite for QuicStream {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<std::result::Result<usize, std::io::Error>> {
        Pin::new(&mut self.get_mut().send)
            .poll_write(cx, buf)
            .map_err(|e| std::io::Error::from(e))
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<std::result::Result<(), std::io::Error>> {
        Pin::new(&mut self.get_mut().send).poll_flush(cx)
    }

    fn poll_shutdown(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<std::result::Result<(), std::io::Error>> {
        Pin::new(&mut self.get_mut().send).poll_shutdown(cx)
    }
}

impl QuicStream {
    /// 发送一个自定义包，一旦发生错误就会关闭流
    pub async fn write_pkg<'b, T>(&mut self, pkg: &T) -> Result<()>
    where
        T: Serializable + SetGetRole + GetAction,
    {
        match split(pkg) {
            Ok(v) => {
                for frame in v {
                    if let Err(e) = self.send.write_all(frame.as_ref()).await {
                        _ = self.close_stream(1).await;
                        return Err(handle_write_error(&e));
                    }
                }
            }
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        }
        if let Err(e) = self.send.flush().await {
            _ = self.close_stream(1).await;
            return Err(Error::msg(e));
        }
        Ok(())
    }

    /// 读取一个数据包元信息，如果发生错误就关闭流
    pub async fn read_metadata(&mut self) -> Result<(PackageMetadata, usize)> {
        let mut buf = [0u8; 3];
        if let Err(e) = self.recv.read_exact(&mut buf).await {
            _ = self.close_stream(1).await;
            return Err(handle_read_exact_error(&e));
        }
        match PackageMetadata::from_sized(buf.to_vec().as_mut()) {
            Ok(e) => Ok(e),
            Err(e) => {
                self.close_stream(1).await;
                Err(e)
            }
        }
    }

    /// 根据数据包元信息读取一个自定义包
    pub async fn read_pkg_with_metadata<T>(
        &mut self,
        metadata: &PackageMetadata,
        pkg: &mut T,
    ) -> Result<()>
    where
        T: Serializable + SetGetRole + GetAction,
    {
        if metadata.get_action() != pkg.get_action() {
            return Err(Error::msg(ERROR_DATA_CHANGED));
        }
        let mut buf = vec![0u8; metadata.len];
        if let Err(e) = self.recv.read_exact(&mut buf).await {
            self.close_stream(1).await;
            return Err(handle_read_exact_error(&e).context("getting package failed"));
        }
        let mut b = match metadata.serialize() {
            Ok(v) => v,
            Err(err) => {
                self.close_stream(1).await;
                return Err(err);
            }
        };
        b.append(&mut buf);
        let mut buffer = Vec::<Vec<u8>>::new();
        buffer.push(b);
        if metadata.has_next {
            loop {
                let (v, md) = match self.read_once().await {
                    Ok(v) => v,
                    Err(e) => {
                        self.close_stream(1).await;
                        return Err(e);
                    }
                };
                buffer.push(v);
                if !md.has_next {
                    break;
                }
            }
        }
        let result = match merge(&buffer) {
            Ok(v) => v,
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        };
        if result.0.pkg_action != pkg.get_action() {
            return Err(Error::msg(ERROR_DATA_CHANGED));
        }
        let size = match pkg.deserialize(result.1.as_slice()) {
            Ok(v) => v,
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        };
        if result.0.len != size {
            self.close_stream(1).await;
            return Err(Error::msg(ERROR_DATA_LEN));
        }
        Ok(())
    }

    pub async fn read_pkg<T>(&mut self, pkg: &mut T) -> Result<()>
    where
        T: Serializable + SetGetRole + GetAction,
    {
        let mut buff = Vec::<Vec<u8>>::new();
        loop {
            match self.read_once().await {
                Ok((v, md)) => {
                    buff.push(v);
                    if !md.has_next {
                        break;
                    }
                }
                Err(e) => {
                    self.close_stream(1).await;
                    return Err(e);
                }
            }
        }
        let p = match merge(&buff) {
            Ok(v) => v,
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        };
        if p.0.pkg_action != pkg.get_action() {
            self.close_stream(1).await;
            return Err(Error::msg(ERROR_DATA_CHANGED));
        }
        let s = match pkg.deserialize(p.1.as_ref()) {
            Ok(v) => v,
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        };
        if s != p.0.len {
            self.close_stream(1).await;
            return Err(Error::msg(ERROR_DATA_LEN));
        }
        Ok(())
    }

    /// 读取一个完整的quic数据包
    async fn read_once(&mut self) -> Result<(Vec<u8>, PackageMetadata)> {
        let mut buf = vec![0u8; 3];
        if let Err(e) = self.recv.read_exact(&mut buf).await {
            self.close_stream(1).await;
            return Err(handle_read_exact_error(&e).context("getting package metadata failed"));
        }
        let (md, _) = match PackageMetadata::from_sized(buf.to_vec().as_ref()) {
            Ok(v) => v,
            Err(e) => {
                self.close_stream(1).await;
                return Err(e);
            }
        };
        let mut buff = vec![0u8; md.len];
        if let Err(e) = self.recv.read_exact(&mut buff).await {
            self.close_stream(1).await;
            return Err(handle_read_exact_error(&e));
        };
        buf.append(&mut buff);
        Ok((buff, md))
    }

    /// 关闭流所属的连接
    pub async fn close_conn(&mut self, code: u32, reason: &str) -> Result<()> {
        self.conn.write().await.close(code.into(), reason);
        Ok(())
    }

    /// 主动关闭流
    pub async fn close_stream(&mut self, code: u32) {
        if self.stat.load(atomic::Ordering::SeqCst) == 0 {
            return;
        }
        match self.send.finish() {
            Ok(()) => {
                _ = self.recv.stop(VarInt::from_u32(code));
            }
            Err(_) => {
                // 流已经关闭了，直接退出即可
            }
        };
        self.conn
            .read()
            .await
            .stream_num
            .fetch_add(-1, atomic::Ordering::SeqCst);
        self.stat.store(0, atomic::Ordering::SeqCst);
    }

    /// 被动等待流关闭
    pub async fn wait_for_close(&mut self) -> Result<Option<VarInt>> {
        if self.stat.load(atomic::Ordering::SeqCst) == 0 {
            return Err(Error::msg("already closed"));
        }
        _ = self.send.finish();
        match self.send.stopped().await {
            Ok(v) => {
                self.conn
                    .read()
                    .await
                    .stream_num
                    .fetch_sub(-1, atomic::Ordering::SeqCst);
                self.stat.store(0, atomic::Ordering::SeqCst);
                Ok(v)
            }
            Err(e) => {
                self.conn
                    .read()
                    .await
                    .stream_num
                    .fetch_add(-1, atomic::Ordering::SeqCst);
                self.stat.store(0, atomic::Ordering::SeqCst);
                Err(handle_stopped_error(&e))
            }
        }
    }

    /// 将流与另外一个双向流进行转发
    pub async fn redir<T>(&mut self, bi_stream: &mut T) -> Result<(u64, u64)>
    where
        T: AsyncRead + AsyncWrite + Unpin,
    {
        match tokio::io::copy_bidirectional(bi_stream, self).await {
            Ok(a) => Ok(a),
            Err(e) => Err(Error::msg(e)),
        }
    }
}

fn handle_connection_error(e: &ConnectionError) -> Error {
    match e {
        ConnectionError::VersionMismatch => {
            Error::msg("ConnectionError: connection version mismatch")
        }
        ConnectionError::TransportError(inner) => Error::msg(format!(
            "ConnectionError: transport error: {}",
            inner.reason
        )),
        ConnectionError::ConnectionClosed(inner) => Error::msg(format!(
            "ConnectionError: connection closed: {}",
            String::from_utf8(inner.reason.to_vec()).unwrap()
        )),
        ConnectionError::ApplicationClosed(inner) => Error::msg(format!(
            "ConnectionError: application closed: {}",
            String::from_utf8(inner.reason.to_vec()).unwrap()
        )),
        ConnectionError::Reset => Error::msg("ConnectionError: reset by peer"),
        ConnectionError::TimedOut => Error::msg("ConnectionError: timed out"),
        ConnectionError::LocallyClosed => Error::msg("ConnectionError: locally closed"),
        ConnectionError::CidsExhausted => Error::msg("ConnectionError: cids exhausted"),
    }
}

fn handle_stopped_error(e: &StoppedError) -> Error {
    match e {
        StoppedError::ConnectionLost(inner) => {
            let err = handle_connection_error(inner);
            err.context("StoppedError: connection lost")
        }
        StoppedError::ZeroRttRejected => Error::msg("StoppedError: ZeroRttRejected"),
    }
}

fn handle_write_error(e: &WriteError) -> Error {
    match e {
        WriteError::Stopped(code) => Error::msg(format!("WriteError: stopped,code is {}", code)),
        WriteError::ConnectionLost(inner) => {
            let err = handle_connection_error(inner);
            err.context("WriteError: connection lost")
        }
        WriteError::ClosedStream => Error::msg("WriteError: connection closed stream"),
        WriteError::ZeroRttRejected => Error::msg("WriteError: ZeroRttRejected"),
    }
}

fn handle_read_exact_error(e: &ReadExactError) -> Error {
    match e {
        ReadExactError::FinishedEarly(size) => {
            Error::msg(format!("ReadExactError: finished early,read {} byte", size))
        }
        ReadExactError::ReadError(inner) => {
            let err = handle_read_error(inner);
            err.context("ReadExactError: read error")
        }
    }
}

fn handle_read_error(e: &ReadError) -> Error {
    match e {
        ReadError::Reset(code) => Error::msg(format!("ReadError: reset by peer, code is {}", code)),
        ReadError::ConnectionLost(inner) => {
            let err = handle_connection_error(inner);
            err.context("ReadError: connection lost")
        }
        ReadError::ClosedStream => Error::msg("ReadError: stream closed"),
        ReadError::IllegalOrderedRead => Error::msg("ReadError: illegal ordered read"),
        ReadError::ZeroRttRejected => Error::msg("ReadError: ZeroRttRejected"),
    }
}
